import numpy as np
import math
import torch
import torch.nn as nn
from torch.autograd import Variable
from Tuatara_Transfromer.utility import clone_layers
from Tuatara_Transfromer.attention import attention


# 我们使用一个类来实现多头注意力机制的处理
class MultiHeadAttention(nn.Module):
    def __init__(self, head_num, embedding_dim, dropout_ratio=0.1):
        """
        :param head_num: head个数，原版论文中为 8
        :param embedding_dim: 表词嵌入的维度，即 d_model
        :param dropout_ratio:
        """

        super(MultiHeadAttention, self).__init__()

        # 测试d_model是否能被 n 整除，
        # 这是因为我们之后要给每个头分配等量的词特征.也就是embedding_dim/head个.
        assert embedding_dim % head_num == 0

        # 每个头获得的词向量维度d_k
        self.d_k = embedding_dim // head_num

        # 传入头数h
        self.head = head_num

        # 然后获得线性层对象，通过nn的Linear实例化，它的内部变换矩阵是embedding_dim x embedding_dim，然后使用clones函数克隆四个，
        # 为什么是四个呢，这是因为在多头注意力中，Q，K，V各需要一个，最后拼接的矩阵还需要一个，因此一共是四个.
        self.linear_layer_list = clone_layers(nn.Linear(embedding_dim, embedding_dim), 4)

        # self.attn为None，它代表最后得到的注意力张量，现在还没有结果所以为None.
        self.attn = None

        # 最后就是一个self.dropout对象，它通过nn中的Dropout实例化而来，置0比率为传进来的参数dropout.
        if dropout_ratio == 0:
            self.dropout_layer = None
        else:
            self.dropout_layer = nn.Dropout(p=dropout_ratio)

    # 前向逻辑函数
    def forward(self, query, key, value, mask_layer=None):
        """
        :param query: Q
        :param key: K
        :param value: V
        :param mask_layer: mask掩码张量
        :return:
        """

        # 如果存在掩码张量mask
        if mask_layer is not None:
            # 使用unsqueeze拓展维度
            mask_layer = mask_layer.unsqueeze(0)

        # batch_size，代表有多少条样本（句子数）.
        batch_size = query.size(0)

        # 之后就进入多头处理环节
        # 首先利用zip将输入QKV与三个线性层组到一起，然后使用for循环，将输入QKV分别传到线性层中，
        # 进行多头拆分：每个头可以代表一部分词特征组成的句子（完整句子，完整词，只是部分维度），
        # 具体使用view方法进行拆分（维度重塑），多加了一个维度h，代表头数，其中的-1代表自适应维度
        # 计算机会根据这种变换自动计算这里的值.然后对第二维和第三维进行转置操作，
        # 为了让代表句子长度维度和词向量维度能够相邻，这样注意力机制才能找到词义与句子位置的关系，
        # 从attention函数中可以看到，利用的是原始输入的倒数第一和第二维.这样我们就得到了每个头的输入.
        query, key, value = \
           [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
            for model, x in zip(self.linear_layer_list, (query, key, value))]

        # 得到每个头的输入后，接下来就是将他们传入到attention中，
        # 这里直接调用我们之前实现的attention函数.同时也将mask和dropout传入其中.
        x, self.attn = attention(query, key, value, mask=mask_layer, dropout=self.dropout_layer)

        # 通过多头注意力计算后，我们就得到了每个头计算结果组成的4维张量，我们需要将其转换为输入的形状以方便后续的计算，
        # 因此这里开始进行第一步处理环节的逆操作，先对第二和第三维进行转置，然后使用contiguous方法，
        # 这个方法的作用就是能够让转置后的张量应用view方法，否则将无法直接使用，
        # 所以，下一步就是使用view重塑形状，变成和输入形状相同.
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)

        # 最后使用线性层列表中的最后一个线性层对输入进行线性变换得到最终的多头注意力结构的输出.
        return self.linear_layer_list[-1](x)


def test():
    # 头数head
    head = 8

    # 词嵌入维度embedding_dim
    embedding_dim = 512

    # 置零比率dropout
    dropout = 0.2




if __name__ == '__main__':

    test()